#ifndef MYAMG_CSRPRECONDKERNEL_CUH
#define MYAMG_CSRPRECONDKERNEL_CUH
#include <cmath>

namespace myamg {
    template<typename Integer, typename Real>
    void prolongationHost(Real *x, const Real *e, const Integer *Pp, const Integer *Pi, const Real *Pv,
                          const size_t nRows) {
#pragma omp parallel for
        for (auto i = 0; i < nRows; i++) {
            Real temp = x[i];
            for (auto j = Pp[i]; j < Pp[i + 1]; j++) {
                temp += Pv[j] * e[Pi[j]];
            }
            x[i] = temp;
        }
    }

    template<typename Integer, typename Real>
    __global__ void prolongationDevice(Real *x, const Real *e, const Integer *Pp, const Integer *Pi, const Real *Pv,
                                       const size_t nRows) {
        const size_t idx = blockIdx.x * blockDim.x + threadIdx.x;
        if (idx < nRows) {
            Real temp = x[idx];
            for (auto j = Pp[idx]; j < Pp[idx + 1]; j++) {
                temp += Pv[j] * e[Pi[j]];
            }
            x[idx] = temp;
        }
    }

    template<typename Integer, typename Real>
    void strengthHost(bool *stg, const Integer *Ap, const Integer *Ai, const Real *Av, const size_t nRows,
                      const double threshold) {
#pragma omp parallel for
        for (auto i = 0; i < nRows; i++) {
            Real theta = 0;
            for (auto j = Ap[i]; j < Ap[i + 1]; j++) {
                if (i != Ai[j]) {
                    theta = std::max(theta, std::abs(Av[j]));
                }
            }
            theta = theta * static_cast<Real>(threshold);
            for (auto j = Ap[i]; j < Ap[i + 1]; j++) {
                stg[j] = (j != i) && (std::abs(Av[j]) > theta);
            }
        }
    }

    template<typename Integer, typename Real>
    __global__ void strengthDevice(bool *stg, const Integer *Ap, const Integer *Ai, const Real *Av, const size_t nRows,
                        const double threshold) {
        const size_t idx = blockIdx.x * blockDim.x + threadIdx.x;
        if (idx < nRows) {
            Real theta = 0;
            for (auto j = Ap[idx]; j < Ap[idx + 1]; j++) {
                if (idx != Ai[j]) {
                    if (std::abs(Av[j]) > theta) {
                        theta = std::abs(Av[j]);
                    }
                }
            }
            theta = theta * static_cast<Real>(threshold);
            for (auto j = Ap[idx]; j < Ap[idx + 1]; j++) {
                stg[j] = (idx != Ai[j]) && (std::abs(Av[j]) > theta);
            }
        }
    }
}

#endif //MYAMG_CSRPRECONDKERNEL_CUH
